(library (dijkstra)
  (export dijkstra-shortest-path)
  (import (except (rnrs base)
                  let-values
                  map)
          (only (guile)
                lambda* λ
                when
                simple-format)
          (ice-9 pretty-print)
          (srfi srfi-1)  ; lists
          (srfi srfi-69)  ; hash tables
          ;; Functional Sets
          (pfds sets)
          ;; Priority Search Queues
          (pfds psqs)
          ;; Bounded Balance Trees
          (pfds bbtrees))


  (define make-empty-set
    (λ (less?)
      (make-set less?)))


  (define set-insert-multiple
    (λ (myset items)
      (cond
       [(null? items) myset]
       [else
        (set-insert-multiple (set-insert myset (car items))
                             (cdr items))])))


  (define set-empty?
    (λ (set)
      (= (set-size set) 0)))


  (define dijkstra-shortest-path
    (lambda* (start-node
              nodes
              get-neighbors
              get-neighbor-distance
              node<
              #:key
              (distance< <))
      "Find the shortest paths between the START-NODE and all other NODES,
given:

GET-NEIGHBORS: A function, that maps from an actual node struct to a list
of node structures.

GET-NEIGHBOR-DISTANCE: A function that returns the distance from one
node to another.

NODE<: A less function, that introduces an order to nodes. This is
important for the underlying functional set implementation.

DISTANCE<: A less function for comparing distances. This is useful for
comparing distances, that are not mere numbers, but might consist of
several attributes. For example how many times a person needs to
change means of transportation to arrive at a destination."
      (define init-unvisited-nodes (set-insert-multiple (make-empty-set node<) nodes))
      (define init-visited-nodes (make-empty-set node<))
      (define init-distances
        (alist->hash-table
         (map (λ (node)
                (if (equal? node start-node)
                    (cons node 0)
                    (cons node +inf.0)))
              nodes)))
      ;; Set distance from start node to itself to 0.
      (define init-routes-table
        (alist->hash-table (list (cons start-node start-node))
                           eq?))

      ;; (hash-table-set! init-distances start-node 0)
      ;; (hash-table-set! init-routes-table start-node start-node)
      ;; Visit an unvisited node with shortest known distance from start
      ;; node. Initially the start node, since all other nodes still have infinite
      ;; distance.
      (let iter ([current-node start-node]
                 [distances° init-distances]
                 [unvisited° init-unvisited-nodes]
                 [visited° init-visited-nodes]
                 [routes° init-routes-table])
        (cond
         ;; Stop, if there are no more unvisited nodes.
         [(set-empty? unvisited°)
          (values distances° routes°)]
         [else
          ;; Calculate distance to every unvisited neighbor from the start
          ;; node. The distance is the distance to current node plus distance to
          ;; the unvisted neighbor).
          (let* ([neighbors (get-neighbors current-node)]
                 ;; Only look at unvisited neighbors.
                 [unvisited-neighbors
                  (filter (λ (neighbor) (set-member? unvisited° neighbor))
                          neighbors)])
            (cond
             ;; If this particular node does not have any unvisited neighbors, go
             ;; back to visiting the next node.
             [(null? unvisited-neighbors)
              ;; Repeat until all nodes visited.
              (iter (set-fold (λ (node acc)
                                (cond [(null? acc) node]
                                      [(distance< (hash-table-ref distances° node)
                                                  (hash-table-ref distances° acc))
                                       node]
                                      [else acc]))
                              '()
                              unvisited°)
                    distances°
                    ;; Mark current node as visited.
                    (set-remove unvisited° current-node)
                    (set-insert visited° current-node)
                    routes°)]
             [else
              ;; Look at the distances to all neighbors and update distances and
              ;; routes accordingly.
              (for-each (λ (neighbor)
                          (let ([start-to-neighbor-distance
                                 (+ (hash-table-ref distances° current-node)
                                    (get-neighbor-distance current-node neighbor))])
                            ;; If distance from the start node to a neighbor node
                            ;; is less than previously known distance for that
                            ;; node, update that distance in the distances
                            ;; table. If a distance is updated, also update what
                            ;; that node's previous node on the path to the node
                            ;; is (the current node).
                            (when (distance< start-to-neighbor-distance
                                             (hash-table-ref distances° neighbor))
                              ;; WARNING: Mutation here! Need
                              ;; something like a functional hash
                              ;; table here ...
                              (hash-table-set! distances°
                                               neighbor start-to-neighbor-distance)
                              (hash-table-set! routes°
                                               neighbor current-node))))
                        unvisited-neighbors)
              ;; Continue with unvisted nodes.
              (iter (set-fold (λ (node acc)
                                (cond [(null? acc) node]
                                      [(distance< (hash-table-ref distances° node)
                                                  (hash-table-ref distances° acc))
                                       node]
                                      [else acc]))
                              '()
                              unvisited°)
                    distances°
                    ;; Mark current node as visited.
                    (set-remove unvisited° current-node)
                    (set-insert visited° current-node)
                    routes°)]))])))))
